// Copyright 2016 The Chromium Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include "media/formats/webm/webm_crypto_helpers.h"

#include "testing/gmock/include/gmock/gmock.h"
#include "testing/gtest/include/gtest/gtest.h"

using ::testing::ElementsAre;

namespace {

const uint8_t kKeyId[] = { 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08 };

} // namespace

namespace media {

inline bool operator==(const SubsampleEntry& lhs, const SubsampleEntry& rhs)
{
    return lhs.clear_bytes == rhs.clear_bytes && lhs.cypher_bytes == rhs.cypher_bytes;
}

TEST(WebMCryptoHelpersTest, EmptyData)
{
    std::unique_ptr<DecryptConfig> decrypt_config;
    int data_offset;
    ASSERT_FALSE(WebMCreateDecryptConfig(nullptr, 0, kKeyId, sizeof(kKeyId),
        &decrypt_config, &data_offset));
}

TEST(WebMCryptoHelpersTest, ClearData)
{
    const uint8_t kData[] = { 0x00, 0x0d, 0x0a, 0x0d, 0x0a };
    std::unique_ptr<DecryptConfig> decrypt_config;
    int data_offset;
    ASSERT_TRUE(WebMCreateDecryptConfig(kData, sizeof(kData), kKeyId,
        sizeof(kKeyId), &decrypt_config,
        &data_offset));
    EXPECT_EQ(1, data_offset);
    EXPECT_FALSE(decrypt_config->is_encrypted());
}

TEST(WebMCryptoHelpersTest, EncryptedButNotEnoughBytes)
{
    const uint8_t kData[] = { 0x01, 0x0d, 0x0a, 0x0d, 0x0a };
    std::unique_ptr<DecryptConfig> decrypt_config;
    int data_offset;
    ASSERT_FALSE(WebMCreateDecryptConfig(kData, sizeof(kData), kKeyId,
        sizeof(kKeyId), &decrypt_config,
        &data_offset));
}

TEST(WebMCryptoHelpersTest, EncryptedNotPartitioned)
{
    const uint8_t kData[] = {
        // Encrypted
        0x01,
        // IV
        0x0d,
        0x0a,
        0x0d,
        0x0a,
        0x0d,
        0x0a,
        0x0d,
        0x0a,
        // Data
        0x01,
        0x02,
    };
    // Extracted from kData and zero extended to 16 bytes.
    const uint8_t kExpectedIv[] = {
        0x0d,
        0x0a,
        0x0d,
        0x0a,
        0x0d,
        0x0a,
        0x0d,
        0x0a,
        0x00,
        0x00,
        0x00,
        0x00,
        0x00,
        0x00,
        0x00,
        0x00,
    };
    std::unique_ptr<DecryptConfig> decrypt_config;
    int data_offset;
    ASSERT_TRUE(WebMCreateDecryptConfig(kData, sizeof(kData), kKeyId,
        sizeof(kKeyId), &decrypt_config,
        &data_offset));
    EXPECT_TRUE(decrypt_config->is_encrypted());
    EXPECT_EQ(std::string(kKeyId, kKeyId + sizeof(kKeyId)),
        decrypt_config->key_id());
    EXPECT_EQ(std::string(kExpectedIv, kExpectedIv + sizeof(kExpectedIv)),
        decrypt_config->iv());
    EXPECT_TRUE(decrypt_config->subsamples().empty());
}

TEST(WebMCryptoHelpersTest, EncryptedPartitionedMissingNumPartitionField)
{
    const uint8_t kData[] = {
        // Encrypted and Partitioned
        0x03,
        // IV
        0x0d,
        0x0a,
        0x0d,
        0x0a,
        0x0d,
        0x0a,
        0x0d,
        0x0a,
    };
    std::unique_ptr<DecryptConfig> decrypt_config;
    int data_offset;
    ASSERT_FALSE(WebMCreateDecryptConfig(kData, sizeof(kData), kKeyId,
        sizeof(kKeyId), &decrypt_config,
        &data_offset));
}

TEST(WebMCryptoHelpersTest, EncryptedPartitionedNotEnoughBytesForOffsets)
{
    const uint8_t kData[] = {
        // Encrypted and Partitioned
        0x03,
        // IV
        0x0d,
        0x0a,
        0x0d,
        0x0a,
        0x0d,
        0x0a,
        0x0d,
        0x0a,
        // Num partitions = 2
        0x02,
        // Partition 0 @ offset 3
        0x00,
        0x00,
        0x00,
        0x03,
    };
    std::unique_ptr<DecryptConfig> decrypt_config;
    int data_offset;
    ASSERT_FALSE(WebMCreateDecryptConfig(kData, sizeof(kData), kKeyId,
        sizeof(kKeyId), &decrypt_config,
        &data_offset));
}

TEST(WebMCryptoHelpersTest, EncryptedPartitionedNotEnoughBytesForData)
{
    const uint8_t kData[] = {
        // Encrypted and Partitioned
        0x03,
        // IV
        0x0d,
        0x0a,
        0x0d,
        0x0a,
        0x0d,
        0x0a,
        0x0d,
        0x0a,
        // Num partitions = 2
        0x02,
        // Partition 0 @ offset 3, partition 2 @ offset 5
        0x00,
        0x00,
        0x00,
        0x03,
        0x00,
        0x00,
        0x00,
        0x05,
        // Should have more than 5 bytes of data
        0x00,
        0x01,
        0x02,
        0x03,
    };
    std::unique_ptr<DecryptConfig> decrypt_config;
    int data_offset;
    ASSERT_FALSE(WebMCreateDecryptConfig(kData, sizeof(kData), kKeyId,
        sizeof(kKeyId), &decrypt_config,
        &data_offset));
}

TEST(WebMCryptoHelpersTest, EncryptedPartitionedNotEnoughBytesForData2)
{
    const uint8_t kData[] = {
        // Encrypted and Partitioned
        0x03,
        // IV
        0x0d,
        0x0a,
        0x0d,
        0x0a,
        0x0d,
        0x0a,
        0x0d,
        0x0a,
        // Num partitions = 2
        0x02,
        // Partition 0 @ offset 3, partition 1 @ offset 5
        0x00,
        0x00,
        0x00,
        0x03,
        0x00,
        0x00,
        0x00,
        0x05,
        // Should have more than 5 bytes of data
        0x00,
        0x01,
        0x02,
        0x03,
        0x04,
    };
    std::unique_ptr<DecryptConfig> decrypt_config;
    int data_offset;
    ASSERT_FALSE(WebMCreateDecryptConfig(kData, sizeof(kData), kKeyId,
        sizeof(kKeyId), &decrypt_config,
        &data_offset));
}

TEST(WebMCryptoHelpersTest, EncryptedPartitionedDecreasingOffsets)
{
    const uint8_t kData[] = {
        // Encrypted and Partitioned
        0x03,
        // IV
        0x0d,
        0x0a,
        0x0d,
        0x0a,
        0x0d,
        0x0a,
        0x0d,
        0x0a,
        // Num partitions = 2
        0x02,
        // Partition 0 @ offset 3, partition 1 @ offset 2
        0x00,
        0x00,
        0x00,
        0x03,
        0x00,
        0x00,
        0x00,
        0x02,
        // Should have more than 5 bytes of data
        0x00,
        0x01,
        0x02,
        0x03,
        0x04,
    };
    std::unique_ptr<DecryptConfig> decrypt_config;
    int data_offset;
    ASSERT_FALSE(WebMCreateDecryptConfig(kData, sizeof(kData), kKeyId,
        sizeof(kKeyId), &decrypt_config,
        &data_offset));
}

TEST(WebMCryptoHelpersTest, EncryptedPartitionedEvenNumberOfPartitions)
{
    const uint8_t kData[] = {
        // Encrypted and Partitioned
        0x03,
        // IV
        0x0d,
        0x0a,
        0x0d,
        0x0a,
        0x0d,
        0x0a,
        0x0d,
        0x0a,
        // Num partitions = 2
        0x02,
        // Partition 0 @ offset 3, partition 1 @ offset 5
        0x00,
        0x00,
        0x00,
        0x03,
        0x00,
        0x00,
        0x00,
        0x05,
        // Should have more than 5 bytes of data
        0x00,
        0x01,
        0x02,
        0x03,
        0x04,
        0x05,
    };
    // Extracted from kData and zero extended to 16 bytes.
    const uint8_t kExpectedIv[] = {
        0x0d,
        0x0a,
        0x0d,
        0x0a,
        0x0d,
        0x0a,
        0x0d,
        0x0a,
        0x00,
        0x00,
        0x00,
        0x00,
        0x00,
        0x00,
        0x00,
        0x00,
    };
    std::unique_ptr<DecryptConfig> decrypt_config;
    int data_offset;
    ASSERT_TRUE(WebMCreateDecryptConfig(kData, sizeof(kData), kKeyId,
        sizeof(kKeyId), &decrypt_config,
        &data_offset));
    EXPECT_TRUE(decrypt_config->is_encrypted());
    EXPECT_EQ(std::string(kKeyId, kKeyId + sizeof(kKeyId)),
        decrypt_config->key_id());
    EXPECT_EQ(std::string(kExpectedIv, kExpectedIv + sizeof(kExpectedIv)),
        decrypt_config->iv());
    EXPECT_THAT(decrypt_config->subsamples(),
        ElementsAre(SubsampleEntry(3, 2), SubsampleEntry(1, 0)));
    EXPECT_EQ(18, data_offset);
}

TEST(WebMCryptoHelpersTest, EncryptedPartitionedOddNumberOfPartitions)
{
    const uint8_t kData[] = {
        // Encrypted and Partitioned
        0x03,
        // IV
        0x0d,
        0x0a,
        0x0d,
        0x0a,
        0x0d,
        0x0a,
        0x0d,
        0x0a,
        // Num partitions = 1
        0x01,
        // Partition 0 @ offset 3,
        0x00,
        0x00,
        0x00,
        0x03,
        // Should have more than 3 bytes of data
        0x00,
        0x01,
        0x02,
        0x03,
        0x04,
        0x05,
    };
    // Extracted from kData and zero extended to 16 bytes.
    const uint8_t kExpectedIv[] = {
        0x0d,
        0x0a,
        0x0d,
        0x0a,
        0x0d,
        0x0a,
        0x0d,
        0x0a,
        0x00,
        0x00,
        0x00,
        0x00,
        0x00,
        0x00,
        0x00,
        0x00,
    };
    std::unique_ptr<DecryptConfig> decrypt_config;
    int data_offset;
    ASSERT_TRUE(WebMCreateDecryptConfig(kData, sizeof(kData), kKeyId,
        sizeof(kKeyId), &decrypt_config,
        &data_offset));
    EXPECT_TRUE(decrypt_config->is_encrypted());
    EXPECT_EQ(std::string(kKeyId, kKeyId + sizeof(kKeyId)),
        decrypt_config->key_id());
    EXPECT_EQ(std::string(kExpectedIv, kExpectedIv + sizeof(kExpectedIv)),
        decrypt_config->iv());
    EXPECT_THAT(decrypt_config->subsamples(), ElementsAre(SubsampleEntry(3, 3)));
    EXPECT_EQ(14, data_offset);
}

} // namespace media
